#!/usr/bin/env python3
"""
重构的BMP数据处理程序
基于test_amp3_signal.py的成功模式，确保能够获取到有效的波形数据
"""

import json
import time
import argparse
import numpy as np
from typing import Dict, List, Optional, Any
import epics


class RefactoredBMPProcessor:
    """重构的BMP数据处理器，采用可靠的PV连接和数据获取策略"""
    
    def __init__(self, config_file: str):
        """初始化BMP处理器"""
        self.config_file = config_file
        self.config = self._load_config()
        self.pvs = {}
        self.connected = False
        
    def _load_config(self) -> Dict[str, Any]:
        """加载配置文件"""
        try:
            with open(self.config_file, 'r', encoding='utf-8') as f:
                config = json.load(f)
            return self._parse_template_config(config)
        except Exception as e:
            print(f"错误: 无法加载配置文件 {self.config_file}: {e}")
            raise
    
    def _parse_template_config(self, template_config: Dict[str, Any]) -> Dict[str, Any]:
        """解析模板配置，生成具体的PV地址"""
        config = {
            'slice_start': template_config.get('slice_start', 7000),
            'slice_end': template_config.get('slice_end', 7050),
            'connection_timeout': template_config.get('connection_timeout', 5.0),
            'chassis': {}
        }
        
        templates = template_config.get('pv_templates', {})
        
        for chassis_info in template_config.get('chassis_list', []):
            chassis_name = chassis_info['name']
            prefix = chassis_info['prefix']
            
            chassis_config = {'bmps': {}}
            
            for bmp_pair in chassis_info.get('bmp_pairs', []):
                bmp1_name = bmp_pair['bmp1']
                bmp2_name = bmp_pair['bmp2']
                
                # 为BMP1生成所有相关的PV地址
                chassis_config['bmps'][bmp1_name] = {
                    'waveforms': self._substitute_template(
                        templates['bmp1_waveforms'], 
                        prefix=prefix, 
                        bmp1=bmp1_name, 
                        bmp2=bmp2_name
                    ),
                    'averages': self._substitute_template(
                        templates['bmp1_averages'], 
                        prefix=prefix, 
                        bmp1=bmp1_name, 
                        bmp2=bmp2_name
                    )
                }
                
                # 为BMP2生成所有相关的PV地址
                chassis_config['bmps'][bmp2_name] = {
                    'waveforms': self._substitute_template(
                        templates['bmp2_waveforms'], 
                        prefix=prefix, 
                        bmp1=bmp1_name, 
                        bmp2=bmp2_name
                    ),
                    'averages': self._substitute_template(
                        templates['bmp2_averages'], 
                        prefix=prefix, 
                        bmp1=bmp1_name, 
                        bmp2=bmp2_name
                    )
                }
            
            config['chassis'][chassis_name] = chassis_config
        
        return config
    
    def _substitute_template(self, template: Dict[str, str], **kwargs) -> Dict[str, str]:
        """替换模板中的占位符"""
        result = {}
        for key, value in template.items():
            result[key] = value.format(**kwargs)
        return result
    
    def connect_pvs(self) -> bool:
        """连接所有PV，采用test_amp3_signal.py的成功模式"""
        print("正在连接PV...")
        
        for chassis_name, chassis_config in self.config['chassis'].items():
            print(f"连接机箱 {chassis_name}...")
            
            for bmp_name, bmp_config in chassis_config['bmps'].items():
                # 连接波形数据PV (采用简单的连接方式)
                for pv_name, pv_addr in bmp_config['waveforms'].items():
                    try:
                        pv = epics.PV(pv_addr)  # 不设置connection_timeout
                        self.pvs[f"{chassis_name}_{bmp_name}_{pv_name}"] = pv
                        print(f"  创建PV: {pv_addr}")
                    except Exception as e:
                        print(f"  ✗ 创建PV失败 {pv_addr}: {e}")
                        return False
                
                # 连接平均值输出PV
                for pv_name, pv_addr in bmp_config['averages'].items():
                    try:
                        pv = epics.PV(pv_addr)
                        self.pvs[f"{chassis_name}_{bmp_name}_{pv_name}_avg"] = pv
                        print(f"  创建PV: {pv_addr}")
                    except Exception as e:
                        print(f"  ✗ 创建PV失败 {pv_addr}: {e}")
                        return False
        
        # 等待连接稳定 (采用test_amp3_signal.py的策略)
        print("等待PV连接稳定...")
        time.sleep(2)
        
        # 检查连接状态
        connected_count = 0
        total_count = len(self.pvs)
        
        for pv_key, pv in self.pvs.items():
            if pv.connected:
                connected_count += 1
                print(f"  ✓ {pv.pvname}")
            else:
                print(f"  ✗ {pv.pvname}")
        
        print(f"PV连接状态: {connected_count}/{total_count}")
        
        if connected_count == total_count:
            self.connected = True
            print("所有PV连接成功！")
            return True
        else:
            print("部分PV连接失败，但继续运行...")
            self.connected = True  # 允许部分连接失败
            return True
    
    def calculate_slice_average(self, waveform_data: Any, start_pos: Optional[int] = None, end_pos: Optional[int] = None) -> float:
        """计算波形数据指定切片的平均值"""
        if waveform_data is None:
            return 0.0
        
        # 如果是标量，直接返回
        if isinstance(waveform_data, (int, float)):
            return float(waveform_data)
        
        # 确保输入数据是numpy数组
        if not isinstance(waveform_data, np.ndarray):
            try:
                waveform_data = np.array(waveform_data)
            except Exception:
                return 0.0
        
        # 检查数组是否为空或者是标量
        if waveform_data.size == 0:
            return 0.0
        
        # 如果是标量，转换为单元素数组
        if waveform_data.ndim == 0:
            return float(waveform_data)
        
        if len(waveform_data) == 0:
            return 0.0
        
        # 确定切片范围
        if start_pos is not None and end_pos is not None:
            start = max(0, start_pos)
            end = min(len(waveform_data), end_pos)
        else:
            start = max(0, self.config['slice_start'])
            end = min(len(waveform_data), self.config['slice_end'])
        
        if start >= end:
            return 0.0
        
        slice_data = waveform_data[start:end]
        return float(np.mean(slice_data))
    
    def get_waveform_data_safely(self, pv_key: str) -> Any:
        """安全地获取波形数据，采用test_amp3_signal.py的策略"""
        if pv_key not in self.pvs:
            return None
        
        pv = self.pvs[pv_key]
        
        if not pv.connected:
            print(f"  警告: PV {pv.pvname} 未连接")
            return None
        
        try:
            # 直接获取数据，不使用额外参数
            data = pv.get()
            
            # 如果获取到的是标量且为0，等待一下再试
            if isinstance(data, (int, float)) and data == 0.0:
                print(f"  获取到标量0值，等待后重试: {pv.pvname}")
                time.sleep(0.5)
                data = pv.get()
            
            return data
            
        except Exception as e:
            print(f"  错误: 获取PV数据失败 {pv.pvname}: {e}")
            return None
    
    def process_single(self) -> Dict[str, Dict[str, float]]:
        """执行单次BMP数据处理"""
        if not self.connected:
            print("错误: PV未连接。请先调用connect_pvs()。")
            return {}

        results = {}
        
        for chassis_name, chassis_config in self.config['chassis'].items():
            results[chassis_name] = {}
            
            for bmp_name, bmp_config in chassis_config['bmps'].items():
                print(f"处理 {chassis_name}/{bmp_name}...")
                
                try:
                    # 直接使用固定的信号段范围 5-15
                    effective_start_pos = 5
                    effective_stop_pos = 15
                    print(f"  使用固定信号段范围: {effective_start_pos} - {effective_stop_pos}")
                    
                    # 获取波形数据 - 从amp波形数据PV获取（修正键名）
                    waveform_va = self.get_waveform_data_safely(f"{chassis_name}_{bmp_name}_va")
                    waveform_vb = self.get_waveform_data_safely(f"{chassis_name}_{bmp_name}_vb")
                    waveform_vc = self.get_waveform_data_safely(f"{chassis_name}_{bmp_name}_vc")
                    waveform_vd = self.get_waveform_data_safely(f"{chassis_name}_{bmp_name}_vd")
                    
                    if waveform_va is None or waveform_vb is None or waveform_vc is None or waveform_vd is None:
                        print(f"  警告: 无法获取波形数据，跳过")
                        results[chassis_name][bmp_name] = {"va_average": 0.0, "vb_average": 0.0, "vc_average": 0.0, "vd_average": 0.0}
                        continue
                    
                    # 打印波形数据信息
                    print(f"  波形数据类型: Va={type(waveform_va)}, Vb={type(waveform_vb)}, Vc={type(waveform_vc)}, Vd={type(waveform_vd)}")
                    
                    # 检查是否获取到有效的数组数据
                    if isinstance(waveform_va, np.ndarray) and len(waveform_va) > 100:
                        print(f"  ✓ 获取到有效波形数据，长度: {len(waveform_va)}")
                        print(f"  数据范围: {np.min(waveform_va):.6f} ~ {np.max(waveform_va):.6f}")
                    else:
                        print(f"  ✗ 波形数据无效: Va={waveform_va}")
                    
                    # 计算背景平均值 (使用配置文件中的背景段范围)
                    background_va = self.calculate_slice_average(waveform_va, self.config['slice_start'], self.config['slice_end'])
                    background_vb = self.calculate_slice_average(waveform_vb, self.config['slice_start'], self.config['slice_end'])
                    background_vc = self.calculate_slice_average(waveform_vc, self.config['slice_start'], self.config['slice_end'])
                    background_vd = self.calculate_slice_average(waveform_vd, self.config['slice_start'], self.config['slice_end'])
                    
                    # 计算信号平均值 (使用有效信号范围)
                    signal_va = self.calculate_slice_average(waveform_va, effective_start_pos, effective_stop_pos)
                    signal_vb = self.calculate_slice_average(waveform_vb, effective_start_pos, effective_stop_pos)
                    signal_vc = self.calculate_slice_average(waveform_vc, effective_start_pos, effective_stop_pos)
                    signal_vd = self.calculate_slice_average(waveform_vd, effective_start_pos, effective_stop_pos)
                    
                    print(f"  背景段范围: {self.config['slice_start']} - {self.config['slice_end']}")
                    print(f"  有效信号段: {effective_start_pos} - {effective_stop_pos}")
                    print(f"  背景平均值 - Va: {background_va:.6f}, Vb: {background_vb:.6f}, Vc: {background_vc:.6f}, Vd: {background_vd:.6f}")
                    print(f"  信号平均值 - Va: {signal_va:.6f}, Vb: {signal_vb:.6f}, Vc: {signal_vc:.6f}, Vd: {signal_vd:.6f}")
                    
                    # 计算背景减除后的BMP电压
                    va_voltage = signal_va - background_va
                    vb_voltage = signal_vb - background_vb
                    vc_voltage = signal_vc - background_vc
                    vd_voltage = signal_vd - background_vd
                    
                    print(f"  背景减除后电压 - Va: {va_voltage:.6f}, Vb: {vb_voltage:.6f}, Vc: {vc_voltage:.6f}, Vd: {vd_voltage:.6f}")
                    
                    # 写入平均值到对应的avg PV
                    try:
                        write_success = 0
                        total_writes = 0
                        
                        # 写入Va平均值
                        va_key = f"{chassis_name}_{bmp_name}_va_avg"
                        if va_key in self.pvs and self.pvs[va_key].connected:
                            self.pvs[va_key].put(va_voltage)
                            print(f"  ✓ 写入 {self.pvs[va_key].pvname}: {va_voltage:.6f}")
                            write_success += 1
                        else:
                            print(f"  ✗ PV {va_key} 未连接或不存在")
                        total_writes += 1
                        
                        # 写入Vb平均值
                        vb_key = f"{chassis_name}_{bmp_name}_vb_avg"
                        if vb_key in self.pvs and self.pvs[vb_key].connected:
                            self.pvs[vb_key].put(vb_voltage)
                            print(f"  ✓ 写入 {self.pvs[vb_key].pvname}: {vb_voltage:.6f}")
                            write_success += 1
                        else:
                            print(f"  ✗ PV {vb_key} 未连接或不存在")
                        total_writes += 1
                        
                        # 写入Vc平均值
                        vc_key = f"{chassis_name}_{bmp_name}_vc_avg"
                        if vc_key in self.pvs and self.pvs[vc_key].connected:
                            self.pvs[vc_key].put(vc_voltage)
                            print(f"  ✓ 写入 {self.pvs[vc_key].pvname}: {vc_voltage:.6f}")
                            write_success += 1
                        else:
                            print(f"  ✗ PV {vc_key} 未连接或不存在")
                        total_writes += 1
                        
                        # 写入Vd平均值
                        vd_key = f"{chassis_name}_{bmp_name}_vd_avg"
                        if vd_key in self.pvs and self.pvs[vd_key].connected:
                            self.pvs[vd_key].put(vd_voltage)
                            print(f"  ✓ 写入 {self.pvs[vd_key].pvname}: {vd_voltage:.6f}")
                            write_success += 1
                        else:
                            print(f"  ✗ PV {vd_key} 未连接或不存在")
                        total_writes += 1
                        
                        print(f"  PV写入完成: {write_success}/{total_writes} 成功")
                        
                        # 回读验证
                        if write_success > 0:
                            print(f"  开始回读验证...")
                            time.sleep(0.1)  # 等待PV更新
                            
                            readback_success = 0
                            tolerance = 1e-4  # 调整容差为更合理的值
                            
                            # 验证Va
                            if va_key in self.pvs and self.pvs[va_key].connected:
                                readback_val = self.pvs[va_key].get()
                                if readback_val is not None:
                                    diff = abs(readback_val - va_voltage)
                                    if diff <= tolerance:
                                        print(f"    ✓ Va回读验证成功: 写入={va_voltage:.6f}, 回读={readback_val:.6f}, 差值={diff:.2e}")
                                        readback_success += 1
                                    else:
                                        print(f"    ✗ Va回读验证失败: 写入={va_voltage:.6f}, 回读={readback_val:.6f}, 差值={diff:.2e}")
                                else:
                                    print(f"    ✗ Va回读失败: 无法获取值")
                            
                            # 验证Vb
                            if vb_key in self.pvs and self.pvs[vb_key].connected:
                                readback_val = self.pvs[vb_key].get()
                                if readback_val is not None:
                                    diff = abs(readback_val - vb_voltage)
                                    if diff <= tolerance:
                                        print(f"    ✓ Vb回读验证成功: 写入={vb_voltage:.6f}, 回读={readback_val:.6f}, 差值={diff:.2e}")
                                        readback_success += 1
                                    else:
                                        print(f"    ✗ Vb回读验证失败: 写入={vb_voltage:.6f}, 回读={readback_val:.6f}, 差值={diff:.2e}")
                                else:
                                    print(f"    ✗ Vb回读失败: 无法获取值")
                            
                            # 验证Vc
                            if vc_key in self.pvs and self.pvs[vc_key].connected:
                                readback_val = self.pvs[vc_key].get()
                                if readback_val is not None:
                                    diff = abs(readback_val - vc_voltage)
                                    if diff <= tolerance:
                                        print(f"    ✓ Vc回读验证成功: 写入={vc_voltage:.6f}, 回读={readback_val:.6f}, 差值={diff:.2e}")
                                        readback_success += 1
                                    else:
                                        print(f"    ✗ Vc回读验证失败: 写入={vc_voltage:.6f}, 回读={readback_val:.6f}, 差值={diff:.2e}")
                                else:
                                    print(f"    ✗ Vc回读失败: 无法获取值")
                            
                            # 验证Vd
                            if vd_key in self.pvs and self.pvs[vd_key].connected:
                                readback_val = self.pvs[vd_key].get()
                                if readback_val is not None:
                                    diff = abs(readback_val - vd_voltage)
                                    if diff <= tolerance:
                                        print(f"    ✓ Vd回读验证成功: 写入={vd_voltage:.6f}, 回读={readback_val:.6f}, 差值={diff:.2e}")
                                        readback_success += 1
                                    else:
                                        print(f"    ✗ Vd回读验证失败: 写入={vd_voltage:.6f}, 回读={readback_val:.6f}, 差值={diff:.2e}")
                                else:
                                    print(f"    ✗ Vd回读失败: 无法获取值")
                            
                            print(f"  回读验证完成: {readback_success}/{write_success} 成功")
                        
                    except Exception as e:
                        print(f"  ✗ 写入avg PV时发生异常: {e}")
                    
                    # 保存结果
                    results[chassis_name][bmp_name] = {
                        "va_average": va_voltage,
                        "vb_average": vb_voltage,
                        "vc_average": vc_voltage,
                        "vd_average": vd_voltage
                    }
                    
                except Exception as e:
                    print(f"  错误: 处理 {chassis_name}/{bmp_name} 时发生异常: {e}")
                    results[chassis_name][bmp_name] = {"va_average": 0.0, "vb_average": 0.0, "vc_average": 0.0, "vd_average": 0.0}
        
        chassis_count = len(results)
        bmp_count = sum(len(bmps) for bmps in results.values())
        print(f"处理完成 - {chassis_count}个机箱, {bmp_count}个BMP")
        
        return results
    
    def run_continuous(self, interval: float = 1.0):
        """连续运行处理"""
        print(f"开始连续处理，间隔: {interval}秒")
        print("按 Ctrl+C 停止...")
        
        try:
            while True:
                start_time = time.time()
                self.process_single()
                
                elapsed = time.time() - start_time
                sleep_time = max(0, interval - elapsed)
                
                if sleep_time > 0:
                    time.sleep(sleep_time)
                    
        except KeyboardInterrupt:
            print("\n程序被用户中断")
        except Exception as e:
            print(f"\n程序运行时发生错误: {e}")
            import traceback
            traceback.print_exc()


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='重构的BMP数据处理程序')
    parser.add_argument('--config', default='config_test.json', help='配置文件路径')
    parser.add_argument('--interval', type=float, default=1.0, help='处理间隔（秒）')
    parser.add_argument('--single', action='store_true', help='执行单次处理后退出')
    
    args = parser.parse_args()
    
    try:
        # 创建处理器
        processor = RefactoredBMPProcessor(args.config)
        
        # 连接PV
        if not processor.connect_pvs():
            print("PV连接失败，退出程序")
            return
        
        if args.single:
            # 执行单次处理
            results = processor.process_single()
            print(f"单次处理结果: {results}")
        else:
            # 运行连续处理
            processor.run_continuous(args.interval)
        
    except Exception as e:
        print(f"程序启动失败: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()